# task4_env.py与task4_train.py是巡航阶段飞机 油门变化 ，输出操作杆指令实现飞行高度变化不大于5米
# 打开DCS之后运行task4_train.py开始训练

import numpy as np
import math
import time
import gym
from gym import spaces
from getDCSdata import get_dcs_data
# 导入手柄控制模块
from vjoy_control import *
import logging
import csv
import os

logger = logging.getLogger(__name__)  # 获取日志记录器

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("results//training.log"),
        logging.StreamHandler()
    ]
)

# 高度偏差阈值
HEIGHT_THRESHOLD = 50
# 高度警戒值为3
HEIGHT_ALERT = 30

# 定义油门变化的间隔步数
THROTTLE_CHANGE_INTERVAL = 100


class ThrottleHeightEnv(gym.Env):
    def __init__(self):
        # 动作空间：操作杆的俯仰和滚转控制，连续动作空间
        self.action_space = spaces.Box(low=-1, high=1, shape=(1,), dtype=np.float32)
        # 观察空间：飞机的高度、速度、姿态等信息，增加油门维度
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(28,), dtype=np.float32)
        self.initial_height = None
        self.current_height = None
        self.flight_data = None
        # 油门初始值
        self.current_throttle = 0.5
        # 记录步数
        self.step_count = 0
        # 记录回合数
        self.episode_count = 0

    def reset(self):
        # 增加回合数
        self.episode_count += 1
        print(f"Starting episode {self.episode_count}")
        # 激活 DCS 窗口
        activate_dcs_window()
        # 获取初始飞行数据
        data = get_dcs_data()
        if data is None:
            logger.warning("Using fallback data after UDP timeout")
            data = self._get_fallback_data()  # 如果获取数据失败，使用模拟数据

        self.initial_height = data["LLA_balt"]
        self.current_height = self.initial_height
        # 模拟飞机处于平飞状态
        # time.sleep(5)
        # 重置步数
        self.step_count = 0
        return self._extract_state(data)

    def step(self, action):
        start_time = time.time()
        # 检查是否需要改变油门
        if self.step_count % THROTTLE_CHANGE_INTERVAL == 0:
            # 简化油门变化规则，减小变化幅度
            throttle_action = np.random.choice([-0.1, 0.1])
            self.current_throttle = np.clip(self.current_throttle + throttle_action, 0, 1)
            if throttle_action > 0:
                throttle_up(throttle_action)
            else:
                throttle_down(-throttle_action)

        # 定义缩放系数
        ACTION_SCALE = 0.3

        # 执行操作杆动作
        # action 是一个单元素数组
        if isinstance(action, np.ndarray):
            if action.size == 1:
                pitch_action = action.item()  # 如果是单元素数组，转换为单个数值
            else:
                raise ValueError("pitch_action 必须是单个数值或单元素数组。")
        else:
            pitch_action = action

        # 对动作值进行缩放
        pitch_action = pitch_action * ACTION_SCALE

        val = abs(pitch_action)
        if pitch_action > 0:
            # 使用vjoy_control中的机头向下（推杆）函数
            pitch_up(val)
        else:
            # 使用vjoy_control中的机头向上（拉杆）函数
            pitch_down(val)

        # 获取新的飞行数据
        data = get_dcs_data()
        self.current_height = data["LLA_balt"]

        # 计算奖励
        reward = self._calculate_reward(data)

        # 计算高度偏差
        height_deviation = abs(self.current_height - self.initial_height)

        # 输出每次训练的奖励
        print(f"Step {self.step_count}: Reward = {reward}, Current Height Deviation = {height_deviation}")

        # 判断是否结束
        done = self._is_done(height_deviation)

        # 记录飞行数据
        self._write_flight_data(data)

        # 增加步数
        self.step_count += 1

        # 获取状态
        state = self._extract_state(data)

        end_time = time.time()
        print(f"Step execution time: {end_time - start_time} seconds")

        return state, reward, done, {}

    def _get_fallback_data(self):
        """当UDP数据不可用时返回模拟数据"""
        return {
            "M_": 0.5, "IAS": 150, "heading_deg": 0,
            "pitch_deg": 5, "bank_deg": 0,
            "Point": [0, 100, 0], "VV": [100, 10, 0],
            "AoA": 5, "ACC": [0, 0, 0],
            "Fuel": 0.8, "Gears": 0, "Flaps": 0,
            "AirBrake": 0, "RPM": [90, 90, 0, 0, 0, 0],
            "LLA_long": 45.019, "LLA_lat": 41.636, "LLA_balt": 3000,
            "Omega": [0.0, 0.0, 0.0]  # 角速度
        }

    def _extract_state(self, data):
        state = [
            data["M_"],
            data["IAS"],
            data["heading_deg"],
            data["pitch_deg"],
            data["bank_deg"],
            *data["Point"],
            *data["VV"],
            data["AoA"],
            *data["ACC"],
            data["Fuel"],
            data["Gears"],
            data["Flaps"],
            data["AirBrake"],
            data["RPM"][0],
            *data["Omega"],
            data["VV_0"],
            data["LLA_long"],
            data["LLA_lat"],
            data["LLA_balt"],
            self.current_throttle  # 加入油门值
        ]
        assert len(state) == 28, f"State length is {len(state)}, expected 28."
        return np.array(state, dtype=np.float32)

    def _is_done(self, height_deviation):
        return height_deviation > HEIGHT_THRESHOLD+500

    def _write_flight_data(self, data):
        csv_filename = "results//flight_data.csv"
        file_exists = os.path.exists(csv_filename)
        with open(csv_filename, 'a', newline='') as csvfile:
            writer = csv.writer(csvfile)
            if not file_exists:
                writer.writerow([
                    "ID", "targetID", "gunID", "flags", "M_", "IAS", "heading_deg", "AoA",
                    "Point_X", "Point_Y", "Point_Z", "VV_X", "VV_Y", "VV_Z", "VV_0",
                    "ACC_X", "ACC_Y", "ACC_Z",
                    "X_Axis_X", "X_Axis_Y", "X_Axis_Z",
                    "Y_Axis_X", "Y_Axis_Y", "Y_Axis_Z",
                    "Z_Axis_X", "Z_Axis_Y", "Z_Axis_Z",
                    "Omega_X", "Omega_Y", "Omega_Z",
                    "pitch_deg", "bank_deg", "yaw_deg",
                    "LLA_long", "LLA_lat", "LLA_balt",
                    "Gears", "Flaps", "AirBrake", "Fuel",
                    "RPM_0", "RPM_1", "RPM_2", "RPM_3", "RPM_4", "RPM_5",
                    "throttle"  # 加入油门列
                ])
            writer.writerow([
                data["ID"], data["targetID"], data["gunID"], data["flags"], data["M_"], data["IAS"],
                data["heading_deg"], data["AoA"],
                data["Point"][0], data["Point"][1], data["Point"][2], data["VV"][0], data["VV"][1], data["VV"][2],
                data["VV_0"],
                data["ACC"][0], data["ACC"][1], data["ACC"][2],
                data["X_Axis"][0], data["X_Axis"][1], data["X_Axis"][2],
                data["Y_Axis"][0], data["Y_Axis"][1], data["Y_Axis"][2],
                data["Z_Axis"][0], data["Z_Axis"][1], data["Z_Axis"][2],
                data["Omega"][0], data["Omega"][1], data["Omega"][2],
                data["pitch_deg"], data["bank_deg"], data["yaw_deg"],
                data["LLA_long"], data["LLA_lat"], data["LLA_balt"],
                data["Gears"], data["Flaps"], data["AirBrake"], data["Fuel"],
                data["RPM"][0], data["RPM"][1], data["RPM"][2], data["RPM"][3], data["RPM"][4], data["RPM"][5],
                self.current_throttle  # 记录油门值
            ])

    def _calculate_reward(self, data):
        # 提取所需数据
        current_height = data["LLA_balt"]
        vertical_speed = data["VV"][1]  # 垂直速度
        pitch_angle = data["pitch_deg"]  # 俯仰角

        # 输出所需信息
        print(f"Current Height: {current_height}, Vertical Speed: {vertical_speed}, Pitch Angle: {pitch_angle}")

        # 计算高度偏差
        height_deviation = abs(current_height - self.initial_height)

        # 权重设置
        weight_height = 0.5
        weight_vertical_speed = 0.2
        weight_pitch_angle = 0.3

        # 高度奖励
        if height_deviation <= 30:
            # 高度偏差在30米之内，给予较多奖励，增长速度较快
            height_reward = 100 - height_deviation ** 2 * (80 / (30 ** 2))
        elif height_deviation <= 50:
            # 高度偏差在30 - 50米之间，给予较少奖励，增长速度较慢
            height_reward = 50 - height_deviation
        else:
            # 高度偏差超过50米，给予惩罚
            height_reward = -0.2 * (height_deviation - HEIGHT_THRESHOLD) ** 2

        # 垂直速度奖励
        target_vertical_speed = 0  # 目标垂直速度
        vertical_speed_error = abs(vertical_speed - target_vertical_speed)
        if vertical_speed_error < 5:
            vertical_speed_reward = 5 * (1 - vertical_speed_error / 5)
        else:
            vertical_speed_reward = -vertical_speed_error

        # 俯仰角奖励
        target_pitch_angle = 0
        pitch_angle_error = abs(pitch_angle - target_pitch_angle)
        if pitch_angle_error <= 1:
            # 1度以内达成目标，给予较高奖励
            pitch_angle_reward = 20 * pitch_angle_error
        elif pitch_angle_error < 5:
            # 1 - 5度以内给予奖励
            pitch_angle_reward = 20 * (1 - (pitch_angle_error - 1) / 4)
        else:
            # 5度以上给予惩罚
            pitch_angle_reward = -pitch_angle_error * 5

        # 加权计算总奖励
        reward = (weight_height * height_reward +
                  weight_vertical_speed * vertical_speed_reward +
                  weight_pitch_angle * pitch_angle_reward)

        return reward